/* Copyright (c) 2025 Huawei Technologies Co., Ltd.
* openUBMC is licensed under Mulan PSL v2.
* You can use this software according to the terms and conditions of the Mulan PSL v2.
* You may obtain a copy of Mulan PSL v2 at:
*          http://license.coscl.org.cn/MulanPSL2
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
* EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
* MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
* See the Mulan PSL v2 for more details.
*/

#define LUA_LIB

#include "luawrapper.h"
#include "luafilter/l_certificate.h"
#include "luafilter/l_measurement_set.h"
#include "luafilter/l_measurement.h"
#include "luafilter/l_signed_measurement.h"
#include "luafilter/l_integrity.h"
#include "luafilter/l_responder_collection.h"
#include "common.h"
#include "spdm_command_impl.h"

ResponderCollection::~ResponderCollection()
{
    Free();
}

Responder *ResponderCollection::GetResponderByName(const string& deviceName)
{
    auto it = find_if(responderCollection.begin(), responderCollection.end(),
        [deviceName](const Responder& obj) {
            return obj.getDeviceName() == deviceName;
    });
    if (it != responderCollection.end()) {
        return &(*it);
    } else {
        return nullptr;
    }
}

void ResponderCollection::Init()
{
    responderCollection = {};

    debug_log(DLOG_NOTICE, "init spdm_lua_lib successful");
}

void ResponderCollection::Free()
{
    for (Responder& obj : responderCollection) {
        obj.ReleaseSpdmCtx();
    }
    responderCollection.clear();
    debug_log(DLOG_NOTICE, "free spdm_lua_lib successful");
    return;
}

int32_t ResponderCollection::GetResponderCount()
{
    return responderCollection.size();
}

void ResponderCollection::NewResponder(const string& deviceName, const string& componentType, const string& inventoryPath)
{
    responderCollection.emplace_back(deviceName, componentType, inventoryPath);
}

int32_t ResponderCollection::NewSpdmCtx(const string& deviceName)
{
    Responder *responder = GetResponderByName(deviceName);
    if (responder == nullptr) {
        debug_log(DLOG_ERROR, "get responder failed, invalid deviceName(%s)", deviceName.c_str());
        return RET_ERR;
    }

    void *context = responder->NewSpdmCtx();
    if (context == nullptr) {
        debug_log(DLOG_ERROR, "execute NewSpdmCtx failed");
        return RET_ERR;
    }
    return RET_OK;
}

int32_t ResponderCollection::InitSpdmCtx(const string& deviceName, uint8_t type)
{
    Responder *responder = GetResponderByName(deviceName);
    if (responder == nullptr) {
        debug_log(DLOG_ERROR, "get responder failed, invalid deviceName(%s)", deviceName.c_str());
        return RET_ERR;
    }

    return responder->InitSpdmCtx(type);
}

int32_t ResponderCollection::CheckSpdmCtx(const string& deviceName)
{
    Responder *responder = GetResponderByName(deviceName);
    if (responder == nullptr) {
        debug_log(DLOG_ERROR, "get responder failed, invalid deviceName(%s)", deviceName.c_str());
        return RET_ERR;
    }

    if (responder->CheckSpdmCtx()) {
        return RET_OK;
    } else {
        return RET_ERR;
    }
}

void ResponderCollection::FreeSpdmCtx(const string& deviceName)
{
    Responder *responder = GetResponderByName(deviceName);
    if (responder == nullptr) {
        debug_log(DLOG_ERROR, "get responder failed, invalid deviceName(%s)", deviceName.c_str());
        return;
    }

    responder->ReleaseSpdmCtx();
}

int32_t ResponderCollection::InitSpdmConnection(const string& deviceName)
{
    Responder *responder = GetResponderByName(deviceName);
    if (responder == nullptr) {
        debug_log(DLOG_ERROR, "get responder failed, invalid deviceName(%s)", deviceName.c_str());
        return RET_ERR;
    }

    return responder->InitConnection();
}

tuple<int32_t, uint8_t> ResponderCollection::GetDigests(const string& deviceName)
{
    Responder *responder = GetResponderByName(deviceName);
    if (responder == nullptr) {
        debug_log(DLOG_ERROR, "get responder failed, invalid deviceName(%s)", deviceName.c_str());
        return {RET_ERR, 0xff};
    }

    uint8_t mask = responder->GetDigests();

    return {RET_OK, mask};
}

luawrap::stack_value ResponderCollection::GetCertificate(lua_State *L, const string& deviceName, uint8_t slot_id)
{
    Responder *responder = GetResponderByName(deviceName);
    if (responder == nullptr) {
        debug_log(DLOG_ERROR, "get responder failed, invalid deviceName(%s)", deviceName.c_str());
        return luaL_error(L, "get certificate failed");
    }

    int32_t res = responder->GetCertificateFromSpdm(slot_id);

    Certificate *cert = responder->GetCertificate(slot_id);
    if (cert == nullptr) {
        debug_log(DLOG_ERROR, "get certificate(%u) failed", slot_id);
        return luaL_error(L, "get certificate failed");
    }
    luawrap::lua_class<LCertificate>::new_object(L, cert);
    return res;
}

int32_t ResponderCollection::ChallengeAuth(const string& deviceName, uint8_t slot_id)
{
    Responder *responder = GetResponderByName(deviceName);
    if (responder == nullptr) {
        debug_log(DLOG_ERROR, "get responder failed, invalid deviceName(%s)", deviceName.c_str());
        return RET_ERR;
    }

    return responder->ChallengeAuth(slot_id);
}

luawrap::stack_value ResponderCollection::GetIntegrity(lua_State *L, const string& deviceName)
{
    Responder *responder = GetResponderByName(deviceName);
    if (responder == nullptr) {
        debug_log(DLOG_ERROR, "get responder failed, invalid deviceName(%s)", deviceName.c_str());
        return luaL_error(L, "get responder failed");
    }

    if (responder->componentIntegrity == nullptr) {
        debug_log(DLOG_ERROR, "get component integrity failed");
        return luaL_error(L, "get component integrity failed");
    }

    ComponentIntegrity *integrity = responder->componentIntegrity;

    luawrap::lua_class<LComponentIntegrity>::new_object(L, integrity->GetType(), integrity->GetVersion(),
        integrity->GetLastUpdated());
    
    return RET_OK;
}

int32_t ResponderCollection::InitMeasurements(const string& deviceName, uint8_t slot_id)
{
    Responder *responder = GetResponderByName(deviceName);
    if (responder == nullptr) {
        debug_log(DLOG_ERROR, "get responder failed, invalid deviceName(%s)", deviceName.c_str());
        return RET_ERR;
    }

    return responder->InitMeasurements(slot_id);
}

luawrap::stack_value ResponderCollection::GetMeasurementSet(lua_State *L, const string& deviceName)
{
    Responder *responder = GetResponderByName(deviceName);
    if (responder == nullptr) {
        debug_log(DLOG_ERROR, "get responder failed, invalid deviceName(%s)", deviceName.c_str());
        return luaL_error(L, "get responder failed");
    }

    if (responder->componentIntegrity == nullptr) {
        debug_log(DLOG_ERROR, "get component integrity failed");
        return luaL_error(L, "get component integrity failed");
    }

    MeasurementSet *set = responder->componentIntegrity->measurementSet;
    if (set == nullptr) {
        debug_log(DLOG_ERROR, "get measurement set failed");
        return RET_ERR;
    }

    luawrap::lua_class<LMeasurementSet>::new_object(L, set->specification, set->summaryType, set->summaryHashAlgorithm,
        set->summaryHash, set->measurementCount);
    
    return RET_OK;
}

luawrap::stack_value ResponderCollection::GetOneMeasurement(lua_State *L, const string& deviceName, uint8_t index)
{
    Responder *responder = GetResponderByName(deviceName);
    if (responder == nullptr) {
        debug_log(DLOG_ERROR, "get responder failed, invalid deviceName(%s)", deviceName.c_str());
        return luaL_error(L, "get responder failed");
    }

    if (responder->componentIntegrity == nullptr) {
        debug_log(DLOG_ERROR, "get component integrity failed");
        return luaL_error(L, "get component integrity failed");
    }

    MeasurementSet *map = responder->componentIntegrity->measurementSet;
    if (map == nullptr) {
        debug_log(DLOG_ERROR, "get measurement set failed");
        return luaL_error(L, "get measurement set failed");
    }

    auto it = map->measurements.find(index);
    if (it == map->measurements.end()) {
        debug_log(DLOG_ERROR, "get measurement failed, invalid index(%u)", index);
        return luaL_error(L, "get measurement failed");
    }

    luawrap::lua_class<LMeasurement>::new_object(L, it->second);
    return RET_OK;
}

luawrap::stack_value ResponderCollection::GetSignedMeasurement(lua_State *L, const string& deviceName, vector<uint8_t> indices,
    uint8_t slot_id, const string& nonce)
{
    Responder *responder = GetResponderByName(deviceName);
    if (responder == nullptr) {
        debug_log(DLOG_ERROR, "get responder failed, invalid deviceName(%s)", deviceName.c_str());
        return luaL_error(L, "get responder failed");
    }

    string version;
    string public_key;
    string hash_algo;
    string sign_algo;
    string measurement;

    tuple<string, string, string, string, string> result;

    try {
        result = responder->GetSignedMeasurement(indices, slot_id, nonce);
    } catch (const runtime_error &e) {
        debug_log(DLOG_ERROR, "caught a error: %s", e.what());
        return luaL_error(L, "get signed measurement failed");
    }
    
    tie(version, public_key, hash_algo, sign_algo, measurement) = result;
    luawrap::lua_class<LSignedMeasurement>::new_object(L, version, hash_algo, sign_algo, public_key, measurement);
    return RET_OK;
}

#ifdef G_ENABLE_TEST
int32_t ResponderCollection::TrySendMsgToCb(const string& deviceName, const string& data)
{
    Responder *responder = GetResponderByName(deviceName);
    if (responder == nullptr) {
        debug_log(DLOG_ERROR, "get responder failed, invalid deviceName(%s)", deviceName.c_str());
        return RET_ERR;
    }

    uint32_t res = responder->msgTransfer->DeviceSendMessage(responder->spdm_context, data.size(), data.data(), 0);
    return res == LIBSPDM_STATUS_SUCCESS ? RET_OK : RET_ERR;
}

luawrap::n_ret ResponderCollection::TryGetMsg(lua_State *L, const string& deviceName)
{
    Responder *responder = GetResponderByName(deviceName);
    if (responder == nullptr) {
        debug_log(DLOG_ERROR, "get responder failed, invalid deviceName(%s)", deviceName.c_str());
        return luaL_error(L, "get responder failed");
    }

    uint8_t *buffer = new uint8_t[SENDER_RECEIVER_MAX_BUFFER_SIZE];
    size_t buffer_len = SENDER_RECEIVER_MAX_BUFFER_SIZE;
    responder->msgTransfer->DeviceReceiveMessage(responder->spdm_context, &buffer_len, (void **)&buffer, 0);
    
    string buffer_str((const char *)buffer, buffer_len);
    delete[] buffer;
    lua_pushstring(L, buffer_str.c_str());
    return 1;
}
#endif // G_ENABLE_TEST

void ResponderCollection::register_to(lua_State *L, luawrap::stack_table &t)
{
    luawrap::restore_stack_top _s(L);
    auto cls = luawrap::lua_class<ResponderCollection>(L);
    cls.ctor<>();
    cls.def("__gc", c_func_wrap(L, ResponderCollection::Free));
    cls.def("free", c_func_wrap(L, ResponderCollection::Free));
    cls.def("init", c_func_wrap(L, ResponderCollection::Init));
    // Responder 管理
    cls.def("get_responder_count", c_func_wrap(L, ResponderCollection::GetResponderCount));
    cls.def("new_responder", c_func_wrap(L, ResponderCollection::NewResponder));
    // 回调注册
    cls.def("set_send_recv_cb", c_func_wrap(L, LCallbacks::set_send_recv_cb));
    // 上下文管理
    cls.def("new_spdm_ctx", c_func_wrap(L, ResponderCollection::NewSpdmCtx));
    cls.def("init_spdm_ctx", c_func_wrap(L, ResponderCollection::InitSpdmCtx));
    cls.def("check_spdm_ctx", c_func_wrap(L, ResponderCollection::CheckSpdmCtx));
    cls.def("free_spdm_ctx", c_func_wrap(L, ResponderCollection::FreeSpdmCtx));
    // 链接管理
    cls.def("init_connection", c_func_wrap(L, ResponderCollection::InitSpdmConnection));
    // AUTH三件套
    cls.def("get_digests", c_func_wrap(L, ResponderCollection::GetDigests));
    cls.def("get_certificate", c_func_wrap(L, ResponderCollection::GetCertificate));
    cls.def("challenge_auth", c_func_wrap(L, ResponderCollection::ChallengeAuth));
    // 完整性报告
    cls.def("get_component_integrity", c_func_wrap(L, ResponderCollection::GetIntegrity));
    // 测量信息管理
    cls.def("init_measurements", c_func_wrap(L, ResponderCollection::InitMeasurements));
    cls.def("get_measurement_set", c_func_wrap(L, ResponderCollection::GetMeasurementSet));
    cls.def("get_one_measurement", c_func_wrap(L, ResponderCollection::GetOneMeasurement));
    cls.def("get_signed_measurement", c_func_wrap(L, ResponderCollection::GetSignedMeasurement));
#ifdef G_ENABLE_TEST
    cls.def("try_send_message_to_cb", c_func_wrap(L, ResponderCollection::TrySendMsgToCb));
    cls.def("try_get_message", c_func_wrap(L, ResponderCollection::TryGetMsg));
#endif // G_ENABLE_TEST

    t.set("core", cls);
}